load("@rules_pkg//:pkg.bzl", "pkg_tar", "pkg_zip")

config_setting(
    name = "windows",
    constraint_values = [
        "@platforms//os:windows",
    ],
)

config_setting(
    name = "rtx_win",
    constraint_values = [
        "@platforms//os:windows",
    ],
    flag_values = {
        "//toolchains/dep_collection:compute_libs": "rtx",
    },
)

pkg_tar(
    name = "include_core",
    package_dir = "include/torch_tensorrt",
    deps = [
        "//core:include",
        "//core/conversion:include",
        "//core/conversion/conversionctx:include",
        "//core/conversion/converters:include",
        "//core/conversion/evaluators:include",
        "//core/conversion/tensorcontainer:include",
        "//core/conversion/var:include",
        "//core/ir:include",
        "//core/lowering:include",
        "//core/lowering/passes:include",
        "//core/partitioning:include",
        "//core/partitioning/partitioningctx:include",
        "//core/partitioning/partitioninginfo:include",
        "//core/partitioning/segmentedblock:include",
        "//core/plugins:impl_include",
        "//core/plugins:include",
        "//core/runtime:include",
        "//core/util:include",
        "//core/util/logging:include",
    ],
)


pkg_tar(
    name = "include_rt",
    package_dir = "include/torch_tensorrt",
    deps = [
        "//core/runtime:include",
    ],
)

pkg_tar(
    name = "include",
    srcs = [
        "//cpp:api_headers",
    ],
    package_dir = "include/torch_tensorrt/",
)

# Non-Windows tar target for lib
pkg_tar(
    name = "lib",
    srcs = [
        "//cpp/lib:libtorchtrt.so",
        "//cpp/lib:libtorchtrt_plugins.so",
        "//cpp/lib:libtorchtrt_runtime.so",
    ],
    mode = "0755",
    package_dir = "lib/",
)

# Non-Windows tar target for lib_rt
pkg_tar(
    name = "lib_rt",
    srcs = ["//cpp/lib:libtorchtrt_runtime.so"],
    mode = "0755",
    package_dir = "lib/",
)

pkg_tar(
    name = "bin",
    srcs = [
        "//cpp/bin/torchtrtc",
    ],
    mode = "0755",
    package_dir = "bin/",
)

# Windows-specific zip target
pkg_zip(
    name = "libtorchtrt_zip",
    srcs = [
        "//core/plugins:impl_include_pkg_files",
        "//core:include_pkg_files",
        "//core/runtime:include_pkg_files",
        "//core/util:include_pkg_files",
        "//core/util/logging:include_pkg_files",
        "//cpp/lib:lib_pkg_files",
        "//:LICENSE",
        "//third_party/torch_tensorrt:BUILD",
    ],
    out = "libtorchtrt.zip",
    package_dir = "torch_tensorrt",
)

# tar target
pkg_tar(
    name = "libtorchtrt_tar",
    srcs = [
        "//:LICENSE",
        "//third_party/torch_tensorrt:BUILD",
    ],
    out = "libtorchtrt.tar.gz",
    extension = "tar.gz",
    package_dir = "torch_tensorrt",
    deps = [
        ":include",
        ":include_core",
        ":lib",
    ] + select({
        ":rtx_win": [],
        ":windows": [],
        "//conditions:default": [":bin"],
    }),
)

# Main target that selects between zip and tar based on platform
alias(
    name = "libtorchtrt",
    actual = select({
        ":rtx_win": ":libtorchtrt_zip",
        ":windows": ":libtorchtrt_zip",
        "//conditions:default": ":libtorchtrt_tar",
    }),
)

# Non-Windows tar target for runtime
pkg_tar(
    name = "libtorchtrt_runtime_tar",
    srcs = [
        "//:LICENSE",
        "//third_party/torch_tensorrt:BUILD",
    ],
    out = "libtorchtrt_runtime.tar.gz",
    extension = "tar.gz",
    package_dir = "torch_tensorrt_runtime",
    deps = [
        ":include_rt",
        ":lib_rt",
    ],
)

# Main target that selects between zip and tar based on platform
alias(
    name = "libtorchtrt_runtime",
    actual = select({
        ":rtx_win": ":libtorchtrt_runtime_zip",
        ":windows": ":libtorchtrt_runtime_zip",
        "//conditions:default": ":libtorchtrt_runtime_tar",
    }),
)
